package com.bzframework.socket.sse;

import com.bzframework.socket.sse.manager.MessagingManager;
import com.bzframework.socket.sse.message.Message;
import com.google.gson.Gson;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.StringUtils;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import reactor.core.Disposable;
import reactor.core.publisher.BufferOverflowStrategy;
import reactor.core.publisher.Mono;
import reactor.core.publisher.SignalType;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * websocket 实现 sse 效果
 * 订阅
 * {
 *  "id": "xxx",
 *  "topic": "/topic",
 *  "type": "sub"
 * }
 * 取消订阅
 * {
 *   "id":"xxx",
 *   "type": "unsub"
 * }
 */
@Slf4j
@RequiredArgsConstructor
public class WebSocketMessagingHandler extends TextWebSocketHandler {

    static Gson gson = new Gson();

    // websocket 当前会话所有订阅信息
    Map<String, Disposable> subs = new ConcurrentHashMap<>();

    private final MessagingManager messagingManager;

    private static final Integer MAX_BUFFER_SIZE = Integer.getInteger("bzframework.websocket.subscriber.buffer-size", 1024);

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {

        // 解析request请求
        MessagingRequest request = gson.fromJson(message.getPayload(), MessagingRequest.class);
        if (request == null) {
            return;
        }
        if (!StringUtils.hasText(request.getId())) {
            sendToSession(session, Message.error(request.getType().name(), null, "id不能为空"));
            return;
        }
        // payload ping -> pong
        if (request.getType() == MessagingRequest.Type.ping) {
            sendToSession(session, Message.pong(request.getId()));
            return;
        }
        if (request.getType() == MessagingRequest.Type.sub) {

            // id 已存在不处理
            Disposable old = subs.get(request.getId());
            if (old != null && !old.isDisposed()) {
                return;
            }

            Disposable subscription = messagingManager
                    .subscribe(request)
                    .onBackpressureBuffer(MAX_BUFFER_SIZE, BufferOverflowStrategy.DROP_OLDEST)
                    .onErrorResume(err -> Mono.just(Message.error(request.getId(), request.getTopic(), err)))
                    .doOnNext(msg -> sendToSession(session, msg))
                    .doFinally(signalType -> {

                        if (signalType == SignalType.ON_COMPLETE) {
                            sendToSession(session, Message.complete(request.getId()));
                        } else {
                            log.debug("{} subscription:{}", signalType, request.getTopic());
                        }

                    })
                    .subscribe();
            if (!subscription.isDisposed()) {
                subs.put(request.getId(), subscription);
            }
        }else if (request.getType() == MessagingRequest.Type.unsub) {
            subs.remove(request.getId()).dispose();
        }else {
            sendToSession(session, Message.error(request.getId(), request.getTopic(), "不支持的类型:" + request.getType()));
        }

    }

    private void sendToSession(WebSocketSession session, Message message) {
        try {
            session.sendMessage(new TextMessage(gson.toJson(message)));
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

}
